'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: loss_zoo.py
@time: 2020-06-05 10:10:09
@desc: 
'''
from ELib.loss.loss import MixSoftmaxCrossEntropyLoss, MixSoftmaxCrossEntropyOHEMLoss
from ELib.loss.loss import LovaszSoftmax, FocalLoss, DiceLoss, ICNetLoss
from ELib.loss.loss import EncNetLoss, PointRendLoss
from jjzhk.config import ZKCFG


def get_segmentation_loss(cfg:ZKCFG):
    model = cfg.MODEL.NAME
    use_ohem = cfg.BASE.OHEM

    kwargs = {'aux': cfg.TRAIN.SOLVER.AUX,
              'aux_weight' : cfg.TRAIN.SOLVER.AUX_WEIGHT,
              'ignore_index' : cfg.TRAIN.SOLVER.IGNORE_INDEX}

    if use_ohem:
        return MixSoftmaxCrossEntropyOHEMLoss(**kwargs)
    elif cfg.BASE.LOSS_NAME == 'lovasz':
        print('Use lovasz loss!')
        return LovaszSoftmax(**kwargs)
    elif cfg.BASE.LOSS_NAME == 'focal':
        print('Use focal loss!')
        return FocalLoss(**kwargs)
    elif cfg.BASE.LOSS_NAME == 'dice':
        print('Use dice loss!')
        return DiceLoss(**kwargs)

    model = model.lower()
    if model == 'icnet':
        print('Use icnet loss!')
        return ICNetLoss(**kwargs)
    elif model == 'encnet':
        print('Use encnet loss!')
        return EncNetLoss(cfg, **kwargs)
    elif model == 'pointrend':
        print('Use pointrend loss!')
        return PointRendLoss(**kwargs)
    else:
        print('Use mix softmax cross entropy loss!')
        return MixSoftmaxCrossEntropyLoss(**kwargs)